import torch
import matplotlib.pyplot as plt

x = torch.linspace(0,100,steps=100).type(torch.FloatTensor)
# print(x.shape)
rand = torch.randn(100) * 10
y = x + rand

x_train = x[:-10]
y_train = y[:-10]

x_test = x[-10:]
y_test = y[-10:]

# print(x.data)
# plt.figure(figsize=(10,8))
# plt.plot(x_train.data.numpy(),y_train.data.numpy(),'o')
# plt.xlabel('x')
# plt.ylabel('y')
# plt.show()

a = torch.rand(1,requires_grad=True)
b = torch.rand(1,requires_grad=True)
learning_rate = 1

for i in range(1000):
    predictions = a.expand_as(x_train) * x_train + b.expand_as(x_train)
    loss = torch.mean((predictions - y_train) ** 2)
    print('loss:',loss.detach().numpy())
    loss.backward()
    a.data.add(-learning_rate * a.grad.data)
    b.data.add(-learning_rate * b.grad.data)
    a.grad.data.zero_()
    b.grad.data.zero_()

x_data = x_train.data.numpy()
plt.figure(figsize=(10,8))
plt.plot(x_train.data.numpy(),y_train.data.numpy(),'o')
plt.plot(x_data,y_train.data.numpy()*x_data + b.data.numpy())
plt.show()